# engine/storage_engine.py
"""
模块职责：作为数据库的存储引擎核心。
它向上层执行器(Executor)提供统一的API，用于操作元组(tuple)和索引。
它将上层的逻辑操作转换为对底层存储管理器(如RecordManager)和物理页面(Page)的调用。
该模块是连接逻辑查询处理和物理数据存储的关键桥梁。
"""

import pickle
from storage.buffer import BufferManager
from .catalog_manager import CatalogManager
from storage.fsm import FreeSpaceManager
from storage.record_manager import RecordManager
from storage.page import PAGE_SIZE
from storage.index_manager import IndexManager

# 定义 RID (Record ID) 为一个元组，包含页面ID和槽ID
RID = tuple[int, int]

class BPlusTreeNode:
    """
    B+树节点在内存中的表示。
    每个节点对象都会被序列化后，完整地存储在数据库的一个页面(Page)中。
    """
    def __init__(self, is_leaf=False, node_size=10):
        self.is_leaf = is_leaf      # bool: 是否为叶子节点
        self.keys = []              # list: 节点的键
        self.pointers = []          # list: 指针。内部节点指向子节点page_id，叶子节点指向数据RID列表
        self.node_size = node_size  # int: 节点的最大容量（键的数量）
        self.parent = None          # int: 父节点的page_id
        self.next_leaf = None       # int: 指向下一个叶子节点的page_id（仅叶子节点有效）

    def is_full(self):
        """检查节点是否已满，需要分裂。"""
        return len(self.keys) >= self.node_size

class StorageEngine:
    """
    存储引擎 (Storage Engine)
    封装了所有物理层面的数据操作，包括元组的增删改查和索引的维护。
    """
    def __init__(self, buffer_manager: BufferManager, catalog_manager: CatalogManager):
        self.buffer_manager = buffer_manager
        self.catalog_manager = catalog_manager
        # 初始化底层的记录管理组件栈
        fsm = FreeSpaceManager(buffer_manager, buffer_manager.file_manager)
        self.record_manager = RecordManager(buffer_manager, fsm, catalog_manager)

    # --- B+树 核心实现 ---

    def _read_node(self, page_id: int) -> BPlusTreeNode:
        """从指定页面读取并反序列化一个B+树节点。"""
        page = self.buffer_manager.get_page(page_id)
        data = page.get_data()
        self.buffer_manager.unpin_page(page_id)
        # 如果页面数据全是0，说明是空页面，返回一个新节点
        if not data.strip(b'\x00'):
            return BPlusTreeNode()
        return pickle.loads(data)

    def _write_node(self, page_id: int, node: BPlusTreeNode):
        """将一个B+树节点序列化并写入到指定页面。"""
        page = self.buffer_manager.get_page(page_id)
        serialized_node = pickle.dumps(node)
        if len(serialized_node) > PAGE_SIZE:
             # 防止节点过大无法存入单个页面
             raise ValueError("B+Tree node exceeds page size.")
        page.set_data(serialized_node)
        self.buffer_manager.unpin_page(page_id, is_dirty=True)

    def _create_new_node(self, is_leaf=False) -> tuple[int, BPlusTreeNode]:
        """分配一个新页面来创建一个新的B+树节点。"""
        new_page = self.buffer_manager.new_page()
        node = BPlusTreeNode(is_leaf=is_leaf)
        # 刚创建的节点是脏的，需要立即写回
        self.buffer_manager.unpin_page(new_page.page_id, is_dirty=True)
        return new_page.page_id, node

    def btree_search(self, root_page_id: int, key: any) -> list[RID]:
        """在B+树中搜索一个键，返回对应的RID列表。"""
        if root_page_id == -1: # 空树
            return []
        node = self._read_node(root_page_id)
        # 1. 从根节点向下遍历，直到找到叶子节点
        while not node.is_leaf:
            i = 0
            while i < len(node.keys) and key >= node.keys[i]:
                i += 1
            child_page_id = node.pointers[i]
            node = self._read_node(child_page_id)
        # 2. 在叶子节点中精确查找键
        try:
            key_index = node.keys.index(key)
            return node.pointers[key_index]
        except ValueError:
            return [] # 未找到

    def btree_insert(self, root_page_id: int, key: any, rid: RID) -> int:
        """B+树插入操作的入口。返回可能已更新的根页面ID。"""
        # 如果是空树，创建根节点并插入
        if root_page_id == -1:
            new_root_page_id, root_node = self._create_new_node(is_leaf=True)
            root_node.keys.append(key)
            root_node.pointers.append([rid])
            self._write_node(new_root_page_id, root_node)
            return new_root_page_id

        # 递归插入
        new_key, new_child_id = self._btree_insert_internal(root_page_id, key, rid)

        # 如果顶层(根节点)发生分裂，需要创建一个新的根节点
        if new_child_id is not None:
            new_root_page_id, new_root_node = self._create_new_node(is_leaf=False)
            new_root_node.keys.append(new_key)
            new_root_node.pointers.extend([root_page_id, new_child_id])
            self._write_node(new_root_page_id, new_root_node)
            return new_root_page_id
        
        return root_page_id

    def _btree_insert_internal(self, page_id: int, key: any, rid: RID) -> tuple[any, int | None]:
        """
        B+树递归插入的内部实现。
        如果当前节点分裂，则返回(分裂产生的键, 新节点的页面ID)以通知父节点。
        如果未分裂，则返回(None, None)。
        """
        node = self._read_node(page_id)
        
        if node.is_leaf:
            # 在叶节点中，如果键已存在，则追加RID；否则插入新的键和RID
            try:
                key_index = node.keys.index(key)
                node.pointers[key_index].append(rid)
            except ValueError:
                insert_pos = self._find_insert_position(node.keys, key)
                node.keys.insert(insert_pos, key)
                node.pointers.insert(insert_pos, [rid])
        else: # 内部节点
            # 找到应该插入的子树，并递归调用
            child_index = self._find_child_index(node.keys, key)
            child_page_id = node.pointers[child_index]
            new_key, new_child_id = self._btree_insert_internal(child_page_id, key, rid)
            
            # 如果子节点发生了分裂，将返回的键和新节点指针插入到当前节点
            if new_child_id is not None:
                insert_pos = self._find_insert_position(node.keys, new_key)
                node.keys.insert(insert_pos, new_key)
                node.pointers.insert(insert_pos + 1, new_child_id)

        # 如果当前节点未满，直接写回并结束
        if not node.is_full():
            self._write_node(page_id, node)
            return None, None

        # --- 节点分裂逻辑 ---
        mid_index = len(node.keys) // 2
        new_node_page_id, new_node = self._create_new_node(is_leaf=node.is_leaf)
        
        if node.is_leaf:
            # 叶子节点分裂：中间的键被复制到新节点，并提升到父节点
            promoted_key = node.keys[mid_index]
            new_node.keys = node.keys[mid_index:]
            new_node.pointers = node.pointers[mid_index:]
            # 更新旧节点
            node.keys = node.keys[:mid_index]
            node.pointers = node.pointers[:mid_index]
        else: # 内部节点
            # 内部节点分裂：中间的键被提升到父节点，不保留在子节点
            promoted_key = node.keys.pop(mid_index)
            new_node.keys = node.keys[mid_index:]
            new_node.pointers = node.pointers[mid_index + 1:]
            # 更新旧节点
            node.keys = node.keys[:mid_index]
            node.pointers = node.pointers[:mid_index + 1]

        self._write_node(page_id, node)
        self._write_node(new_node_page_id, new_node)
        # 返回被提升的键和新节点的ID，通知上层节点进行更新
        return promoted_key, new_node_page_id
    
    def _find_insert_position(self, keys, key):
        """找到新键在有序列表中的插入位置。"""
        pos = 0
        while pos < len(keys) and keys[pos] < key:
            pos += 1
        return pos

    def _find_child_index(self, keys, key):
        """在内部节点中，根据键找到对应的子节点指针的索引。"""
        idx = 0
        while idx < len(keys) and key >= keys[idx]:
            idx += 1
        return idx

    # --- 元组操作API (全部委托给RecordManager) ---

    def insert_tuple(self, table_name: str, values: list) -> RID:
        """Insert tuple via RecordManager and update indexes."""
        rid = self.record_manager.insert_record(table_name, tuple(values))
        # Ensure catalog page list
        page_id, _ = rid
        schema = self.catalog_manager.get_table_schema(table_name)
        if schema and page_id not in schema.pages:
            self.catalog_manager.add_page_to_table(table_name, page_id)
        # Update indexes using B+Tree
        if schema:
            for col_name, idx_schema in schema.indexes.items():
                try:
                    col_idx = schema.column_order.index(col_name)
                    key = values[col_idx]
                except ValueError:
                    continue
                im = IndexManager(self.buffer_manager, idx_schema.root_page_id)
                im.insert(int(key), rid)  # assume INT key
                if idx_schema.root_page_id is None:
                    idx_schema.root_page_id = im.root_page_id
                    # save catalog with new root
                    self.catalog_manager.save_catalog()
        return rid

    def get_tuple(self, table_name: str, rid: RID) -> list | None:
        """根据RID获取一个元组。实际工作由RecordManager完成。"""
        record = self.record_manager.get_record(table_name, rid)
        return list(record) if record else None

    def delete_tuple(self, table_name: str, rid: RID):
        """根据RID删除一个元组。实际工作由RecordManager完成。"""
        # 注意：删除元组后，索引没有更新。在一个完整的DBMS中，这里也需要更新索引。
        self.record_manager.delete_record(table_name, rid)

    def update_tuple(self, table_name: str, rid: RID, new_tuple: list):
        """更新一个元组。采用“先删除后插入”的策略，以简化实现。"""
        # 注意：这是一个简化的实现。高效的更新应该是“就地更新”(in-place update)。
        # 同样，这里也需要处理索引的更新。
        self.record_manager.delete_record(table_name, rid)
        self.record_manager.insert_record(table_name, tuple(new_tuple))

    def scan(self, table_name: str):
        """对表进行全表扫描。实际工作由RecordManager完成。"""
        for _, record in self.record_manager.scan_with_rid(table_name):
            yield record

    def scan_with_rid(self, table_name: str):
        """带RID的全表扫描。实际工作由RecordManager完成。"""
        return self.record_manager.scan_with_rid(table_name)

    # --- 辅助方法 ---

    def _update_indexes_for_insert(self, schema, values: list, rid: RID):
        """当插入一条新记录时，更新该表上定义的所有索引。"""
        if not schema.indexes: return
        for col_name, index_schema in schema.indexes.items():
            try:
                col_index = schema.column_order.index(col_name)
                key = values[col_index]
                # 获取索引的根页面ID
                root_page_id = self.catalog_manager.get_index_root_page_id(index_schema.index_name)
                # 将（键，RID）插入B+树
                new_root_page_id = self.btree_insert(root_page_id, str(key), rid)
                # 如果B+树的根节点发生变化，更新目录中的元数据
                if new_root_page_id != root_page_id:
                    self.catalog_manager.update_index_root_page_id(index_schema.index_name, new_root_page_id)
            except (ValueError, IndexError):
                # 如果找不到列或值，则跳过该索引
                continue
    
    def find_table_by_page(self, page_id: int) -> str:
        """根据页面ID反向查找它属于哪个表。"""
        all_tables = self.catalog_manager._catalog.tables
        for table_name, schema in all_tables.items():
            if page_id in schema.pages:
                return table_name
        raise RuntimeError(f"无法找到页面 {page_id} 所属的表。")